import math
import matplotlib.pyplot as plt
import numpy as np
from typing import List, Tuple
import datetime
import os


# GEOPro，使用球面距离计算城市之间的距离
class GEOPro:
    def __init__(self, num_city: int, data: np.ndarray, is_geo_coordinate: bool = False):
        # 基础参数
        self.m = 50  # 蚂蚁数量
        self.alpha = 1.0  # 信息素重要程度因子
        self.beta = 5.0  # 启发函数重要因子
        self.rho = 0.1  # 信息素挥发因子
        self.Q = 1  # 常量系数
        self.elite_weight = 2.0  # 精英蚂蚁信息素权重

        # 自适应参数
        self.alpha_max = 2.0
        self.alpha_min = 0.5
        self.beta_max = 8.0
        self.beta_min = 2.0

        self.num_city = num_city
        self.is_geo_coordinate = is_geo_coordinate
        self.location = data

        if self.is_geo_coordinate:
            self.location = self.convert_to_radians(data)

        self.Tau = np.ones([num_city, num_city]) * 0.1
        self.Table = [[0 for _ in range(num_city)] for _ in range(self.m)]
        self.iter = 1
        self.iter_max = 3
        self.dis_mat = self.compute_dis_mat(num_city, self.location)
        self.Eta = 10. / (self.dis_mat + 1e-10)
        self.paths = None

        self.iter_x = []
        self.iter_y = []

        self.best_path_history = []
        self.best_length_history = []

        self.local_search_prob = 0.3  # 进行局部搜索的概率

        self.greedy_init(self.dis_mat, min(50, num_city), num_city)

    def convert_to_radians(self, coordinates: np.ndarray) -> np.ndarray:
        """
        将地理坐标(经纬度)转换为弧度表示
        公式: rad = PI * (deg + 5.0 * min / 3.0) / 180.0
        """
        PI = 3.141592
        radians_coords = np.zeros_like(coordinates)

        for i in range(len(coordinates)):
            for j in range(2):
                x = coordinates[i, j]
                deg = int(x)
                minutes = x - deg
                radians_coords[i, j] = PI * (deg + 5.0 * minutes / 3.0) / 180.0

        return radians_coords

    def geo_distance(self, coord1: np.ndarray, coord2: np.ndarray) -> float:
        """
        计算两个地理坐标点之间的距离
        使用地理距离公式:
        计算公式来源:http://comopt.ifi.uni-heidelberg.de/software/TSPLIB95/TSPFAQ.html

        For converting coordinate input to longitude and latitude in radian:

        PI = 3.141592;
        deg = (int) x[i];
        min = x[i]- deg;
        rad = PI * (deg + 5.0 * min/ 3.0) / 180.0;

        For computing the geographical distance:
        RRR = 6378.388;
        q1 = cos( longitude[i] - longitude[j] );
        q2 = cos( latitude[i] - latitude[j] );
        q3 = cos( latitude[i] + latitude[j] );
        dij = (int) ( RRR * acos( 0.5*((1.0+q1)*q2 - (1.0-q1)*q3) ) + 1.0);
        """
        RRR = 6378.388
        longitude_i, latitude_i = coord1[0], coord1[1]
        longitude_j, latitude_j = coord2[0], coord2[1]

        q1 = math.cos(longitude_i - longitude_j)
        q2 = math.cos(latitude_i - latitude_j)
        q3 = math.cos(latitude_i + latitude_j)

        inner = 0.5 * ((1.0 + q1) * q2 - (1.0 - q1) * q3)
        inner = max(-1.0, min(1.0, inner))

        dij = int(RRR * math.acos(inner) + 1.0)
        return dij

    def compute_dis_mat(self, num_city: int, location: np.ndarray) -> np.ndarray:
        """
        计算城市之间的距离矩阵
        """
        dis_mat = np.zeros((num_city, num_city))
        for i in range(num_city):
            for j in range(num_city):
                if i == j:
                    dis_mat[i][j] = np.inf
                    continue
                if self.is_geo_coordinate:
                    dis_mat[i][j] = self.geo_distance(location[i], location[j])
                else:
                    a = location[i]
                    b = location[j]
                    tmp = np.sqrt(sum([(x[0] - x[1]) ** 2 for x in zip(a, b)]))
                    dis_mat[i][j] = tmp
        return dis_mat

    def greedy_init(self, dis_mat: np.ndarray, num_total: int, num_city: int):
        start_index = 0
        result = []
        for i in range(num_total):
            rest = [x for x in range(0, num_city)]
            if start_index >= num_city:
                start_index = np.random.randint(0, num_city)
                result.append(result[start_index].copy())
                continue
            current = start_index
            rest.remove(current)
            result_one = [current]
            while len(rest) != 0:
                tmp_min = math.inf
                tmp_choose = -1
                for x in rest:
                    if dis_mat[current][x] < tmp_min:
                        tmp_min = dis_mat[current][x]
                        tmp_choose = x

                current = tmp_choose
                result_one.append(tmp_choose)
                rest.remove(tmp_choose)
            result.append(result_one)
            start_index += 1
        pathlens = self.compute_paths(result)
        sortindex = np.argsort(pathlens)
        index = sortindex[0]
        result = result[index]
        for i in range(len(result) - 1):
            s = result[i]
            s2 = result[i + 1]
            self.Tau[s][s2] += 1
        self.Tau[result[-1]][result[0]] += 1

    # 轮盘赌选择
    def rand_choose(self, p: List[float]) -> int:
        x = np.random.rand()
        for i, t in enumerate(p):
            x -= t
            if x <= 0:
                break
        return i

    def get_ants(self, num_city: int):
        """
        生成蚁群
        """
        for i in range(self.m):
            start = np.random.randint(num_city)
            self.Table[i][0] = start
            unvisit = list([x for x in range(num_city) if x != start])
            current = start
            j = 1
            while len(unvisit) != 0:
                P = []
                for v in unvisit:
                    P.append(self.Tau[current][v] ** self.alpha * self.Eta[current][v] ** self.beta)
                P_sum = sum(P)
                if P_sum == 0:
                    index = np.random.randint(len(unvisit))
                else:
                    P = [x / P_sum for x in P]
                    index = self.rand_choose(P)
                current = unvisit[index]
                self.Table[i][j] = current
                unvisit.remove(current)
                j += 1

            if np.random.rand() < self.local_search_prob:
                self.local_optimization(i)

    def local_optimization(self, ant_index: int):
        """
        对指定蚂蚁的路径进行2-opt局部优化
        """
        path = self.Table[ant_index]
        best_path = path[:]
        best_distance = self.compute_pathlen(path, self.dis_mat)

        improved = True
        while improved:
            improved = False
            for i in range(1, len(path) - 2):
                for j in range(i + 1, len(path)):
                    if j - i == 1:
                        continue
                    new_path = path[:]
                    new_path[i:j] = path[i:j][::-1]
                    new_distance = self.compute_pathlen(new_path, self.dis_mat)
                    if new_distance < best_distance:
                        best_path = new_path[:]
                        best_distance = new_distance
                        improved = True
            path = best_path[:]

        self.Table[ant_index] = best_path

    def compute_pathlen(self, path: List[int], dis_mat: np.ndarray) -> float:
        """
        计算一条路径的长度
        """
        a = path[0]
        b = path[-1]
        result = dis_mat[a][b]
        for i in range(len(path) - 1):
            a = path[i]
            b = path[i + 1]
            result += dis_mat[a][b]
        return result

    def compute_paths(self, paths: List[List[int]]) -> List[float]:
        """
        计算一个群体的长度
        """
        result = []
        for one in paths:
            length = self.compute_pathlen(one, self.dis_mat)
            result.append(length)
        return result

    def update_Tau(self, best_path: List[int], best_length: float):
        """
        更新信息素
        """
        delta_tau = np.zeros([self.num_city, self.num_city])
        paths = self.compute_paths(self.Table)

        for i in range(self.m):
            for j in range(self.num_city - 1):
                a = self.Table[i][j]
                b = self.Table[i][j + 1]
                delta_tau[a][b] += self.Q / paths[i]
                delta_tau[b][a] += self.Q / paths[i]
            a = self.Table[i][0]
            b = self.Table[i][-1]
            delta_tau[a][b] += self.Q / paths[i]
            delta_tau[b][a] += self.Q / paths[i]

        if best_path is not None:
            for j in range(self.num_city - 1):
                a = best_path[j]
                b = best_path[j + 1]
                delta_tau[a][b] += self.elite_weight * self.Q / best_length
                delta_tau[b][a] += self.elite_weight * self.Q / best_length
            a = best_path[0]
            b = best_path[-1]
            delta_tau[a][b] += self.elite_weight * self.Q / best_length
            delta_tau[b][a] += self.elite_weight * self.Q / best_length

        self.Tau = (1 - self.rho) * self.Tau + delta_tau

    def adaptive_parameters(self, iter_ratio: float):
        """
        自适应调整参数
        """
        self.alpha = self.alpha_max - (self.alpha_max - self.alpha_min) * iter_ratio
        self.beta = self.beta_min + (self.beta_max - self.beta_min) * iter_ratio

    def aco(self) -> Tuple[float, List[int]]:

        best_length = math.inf
        best_path = None

        for cnt in range(self.iter_max):
            # 自适应调整参数
            iter_ratio = cnt / self.iter_max
            self.adaptive_parameters(iter_ratio)

            self.get_ants(self.num_city)  # out>>self.Table
            self.paths = self.compute_paths(self.Table)

            tmp_length = min(self.paths)
            tmp_path = self.Table[self.paths.index(tmp_length)]

            if tmp_length < best_length:
                best_length = tmp_length
                best_path = tmp_path[:]

                print(f"*** 新最优解 *** Iteration {cnt}: best length = {best_length:.2f}")

            self.update_Tau(best_path, best_length)

            self.iter_x.append(cnt)
            self.iter_y.append(best_length)

            self.best_path_history.append(best_path[:])
            self.best_length_history.append(best_length)

            if cnt % 10 == 0:
                print(f"Iteration {cnt}: best length = {best_length:.2f}, current generation best = {tmp_length:.2f}")

        return best_length, best_path

    def save_tour_file(self, best_path: List[int], best_length: float):
        """
        保存结果为 tour 文件格式
        """
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"GEOPro_result_{timestamp}.tour"

        with open(filename, 'w') as f:
            f.write(f"NAME : GEOPro_result\n")
            f.write(f"COMMENT : Optimum solution found by GEOPro\n")
            f.write(f"COMMENT : Solution length: {best_length:.2f}\n")
            f.write(f"COMMENT : Generated at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("TYPE : TOUR\n")
            f.write(f"DIMENSION : {self.num_city}\n")
            f.write("TOUR_SECTION\n")

            for city_index in best_path:
                f.write(f"{city_index + 1}\n")
            f.write("-1\n")
            f.write("EOF\n")

        print(f"结果已保存至: {filename}")

    def run(self) -> Tuple[np.ndarray, float, List[int]]:

        print("开始运行改进蚁群算法...（使用球面距离）")
        print(f"城市数量: {self.num_city}, 最大迭代次数: {self.iter_max}")
        print(f"蚂蚁数量: {self.m}, 参数设置: alpha={self.alpha}, beta={self.beta}, rho={self.rho}")

        best_length, best_path = self.aco()

        print("=" * 50)
        print("算法运行完成!")
        print(f"最终最优路径长度: {best_length:.2f}")
        print("最佳路径节点序列:")
        print(" -> ".join(map(str, best_path)))
        print("=" * 50)

        self.save_tour_file(best_path, best_length)

        if self.is_geo_coordinate:
            return self.convert_to_radians(self.location[best_path]), best_length, best_path
        else:
            return self.location[best_path], best_length, best_path


def read_tsp(path: str) -> List[List[float]]:
    lines = open(path, 'r').readlines()
    assert 'NODE_COORD_SECTION\n' in lines
    index = lines.index('NODE_COORD_SECTION\n')
    data = lines[index + 1:-1]
    tmp = []
    for line in data:
        line = line.strip().split(' ')
        if line[0] == 'EOF':
            continue
        tmpline = []
        for x in line:
            if x == '':
                continue
            else:
                tmpline.append(float(x))
        if tmpline == []:
            continue
        tmp.append(tmpline)
    data = tmp
    return data


if __name__ == "__main__":

    data = read_tsp('../data/att48.tsp')
    data = np.array(data)
    data = data[:, 1:]

    improved_aco = GEOPro(num_city=data.shape[0], data=data.copy(), is_geo_coordinate=True)

    Best_path_coords, Best_length, Best_path_nodes = improved_aco.run()

    Best_path_coords_vstacked = np.vstack([Best_path_coords, Best_path_coords[0]])
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(Best_path_coords_vstacked[:, 0], Best_path_coords_vstacked[:, 1], 'o-', markersize=4)

    for i, (x, y) in enumerate(Best_path_coords):
        plt.annotate(str(Best_path_nodes[i]), (x, y), xytext=(3, 3), textcoords='offset points', fontsize=8)

    plt.title('GEOPro')
    plt.xlabel('X')
    plt.ylabel('Y')

    plt.subplot(1, 2, 2)
    plt.plot(improved_aco.iter_x, improved_aco.iter_y)
    plt.title('Convergence Curve')
    plt.xlabel('Iteration')
    plt.ylabel('Best Length')

    plt.tight_layout()
    plt.show()

